# -*- coding: utf-8 -*-
import os
import psutil

from utils.log_utils import logger


class LocustLock:
    """用于防止重复运行 Locust 的跨进程文件锁管理器"""

    def __init__(self, lock_file: str = "locust.lock"):
        self.lock_file = lock_file

    def is_locked(self) -> bool:
        """判断是否已有任务运行中"""
        if not os.path.exists(self.lock_file):
            return False
        try:
            with open(self.lock_file, "r") as f:
                pid = int(f.read().strip())
            logger.info(f"[锁检查] 当前锁定PID: {pid}")
            if psutil.pid_exists(pid):
                return True
            else:
                os.remove(self.lock_file)  # 清除无效锁
                return False
        except Exception as e:
            logger.error(f"[锁异常] is_locked 错误: {e}")
            if os.path.exists(self.lock_file):
                os.remove(self.lock_file)
            return False

    def acquire(self, pid: int):
        """创建锁"""
        with open(self.lock_file, "w") as f:
            f.write(str(pid))
        logger.info(f"[锁创建] 进程 {pid} 获取锁成功")

    def release(self):
        """释放锁"""
        if os.path.exists(self.lock_file):
            os.remove(self.lock_file)
            logger.info(f"[锁释放] 已删除锁文件")
